import os
import pickle
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import classification_report, accuracy_score
import pandas as pd

# 加载数据集
data = pd.read_csv('Normalization_total_extend.csv')

# 假设最后一列是已经编码好的类别标签
X = data.iloc[:, :-1]  # 特征
y = data.iloc[:, -1]  # 类别标签（已编码）

# 划分数据集为训练集和测试集，比例为8:2
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 定义模型文件路径
model_file = 'NaiveBayesModel/naive_bayes_model.pkl'

# 检查模型文件是否存在
if os.path.exists(model_file):
    # 加载模型
    with open(model_file, 'rb') as f:
        nb_model = pickle.load(f)
        print("Model loaded from file.")
else:
    # 初始化朴素贝叶斯分类器（假设特征是连续的，使用高斯朴素贝叶斯）
    nb_model = GaussianNB()

    # 训练模型
    nb_model.fit(X_train, y_train)

    # 保存模型
    with open(model_file, 'wb') as f:
        pickle.dump(nb_model, f)
        print("Model saved to file.")

    # 预测测试集
y_pred = nb_model.predict(X_test)

# 性能评估
accuracy = accuracy_score(y_test, y_pred)
print(f'Accuracy: {accuracy:.4f}')

# 由于标签已经编码，直接使用y_test和y_pred计算precision, recall, f1-score
report = classification_report(y_test, y_pred, target_names=None)  # target_names可能无法直接使用，除非你有原始的标签名称列表
print(report)